# -*- coding: utf-8 -*-

import torch

from data.build_dataset import build_dataset
from data.build_transforms import build_transforms


def make_data_sampler(dataset, phase="train"):
    assert phase in ("train", "test")

    if phase == "train":
        sampler = torch.utils.data.sampler.RandomSampler(dataset)
    else:
        sampler = torch.utils.data.sampler.SequentialSampler(dataset)
    return sampler


def make_batch_data_sampler(sampler, batch_size):
    batch_sampler = torch.utils.data.sampler.BatchSampler(
        sampler, batch_size, drop_last=False
    )
    return batch_sampler


def build_data_loader(cfg):
    phase = cfg.PHASE
    transforms = build_transforms(cfg)
    dataset = build_dataset(cfg, transforms)
    data_sampler = make_data_sampler(dataset, phase=phase)
    batch_sampler = make_batch_data_sampler(data_sampler,
                                            batch_size=cfg.DATALOADER.BATCH_SIZE)

    data_loader = torch.utils.data.DataLoader(dataset,
                                              num_workers=cfg.DATALOADER.NUM_WORKERS,
                                              batch_sampler=batch_sampler)
    return data_loader
